#include <osgDB/ReadFile>
#include <osgUtil/Optimizer>
#include <osg/ComputeBoundsVisitor>
#include <osgViewer/Viewer>
#include <osg/Camera>
#include <osg/Texture2D>
#include <osg/Image>
#include <osgDB/WriteFile>
#include <osg/GraphicsContext>
#include <osg/Material>
#include <osg/PolygonMode>
#include <osg/LightModel>
#include <string>
#include <filesystem>
#include <iostream>
#include <iomanip>

struct RenderParameters {
    int imageWidth = 1024;
    int imageHeight = 1024;
    float marginRatio = 0.1f;     // Margin ratio, default is 10% of max size
    float viewDistance = 2.0f;    // View distance coefficient, default is 2x max size
    int frameCount = 3;           // Number of frames to render for each view
    std::string outputDir = ".";  // Output directory
};

void printHelp(const char* programName)
{
    std::cout << "Usage: " << programName << " <osgb file path> [options]" << std::endl;
    std::cout << "\nOptions:" << std::endl;
    std::cout << std::left << std::setw(30) << "  -w, --width <value>"
        << "Output image width (default: 1024)" << std::endl;
    std::cout << std::left << std::setw(30) << "  -h, --height <value>"
        << "Output image height (default: 1024)" << std::endl;
    std::cout << std::left << std::setw(30) << "  -m, --margin <value>"
        << "Margin ratio (default: 0.1)" << std::endl;
    std::cout << std::left << std::setw(30) << "  -d, --distance <value>"
        << "View distance coefficient (default: 2.0)" << std::endl;
    std::cout << std::left << std::setw(30) << "  -f, --frames <value>"
        << "Number of frames to render (default: 3)" << std::endl;
    std::cout << std::left << std::setw(30) << "  -o, --output <path>"
        << "Output directory path (default: current directory)" << std::endl;
    std::cout << std::left << std::setw(30) << "  --help"
        << "Display this help information" << std::endl;
}

bool parseCommandLine(int argc, char** argv, std::string& inputFile, RenderParameters& params)
{
    if (argc < 2)
    {
        printHelp(argv[0]);
        return false;
    }

    inputFile = argv[1];

    for (int i = 2; i < argc; i++)
    {
        std::string arg = argv[i];

        if (arg == "--help") {
            printHelp(argv[0]);
            return false;
        }

        if (i + 1 >= argc) {
            std::cout << "Error: " << arg << " requires a parameter value" << std::endl;
            return false;
        }

        if (arg == "-w" || arg == "--width") {
            params.imageWidth = std::stoi(argv[++i]);
        }
        else if (arg == "-h" || arg == "--height") {
            params.imageHeight = std::stoi(argv[++i]);
        }
        else if (arg == "-m" || arg == "--margin") {
            params.marginRatio = std::stof(argv[++i]);
        }
        else if (arg == "-d" || arg == "--distance") {
            params.viewDistance = std::stof(argv[++i]);
        }
        else if (arg == "-f" || arg == "--frames") {
            params.frameCount = std::stoi(argv[++i]);
        }
        else if (arg == "-o" || arg == "--output") {
            params.outputDir = argv[++i];
        }
        else {
            std::cout << "Unknown parameter: " << arg << std::endl;
            return false;
        }
    }

    // Validate parameters
    if (params.imageWidth <= 0 || params.imageHeight <= 0) {
        std::cout << "Error: Image dimensions must be greater than 0" << std::endl;
        return false;
    }
    if (params.marginRatio < 0) {
        std::cout << "Error: Margin ratio must be greater than or equal to 0" << std::endl;
        return false;
    }
    if (params.viewDistance <= 0) {
        std::cout << "Error: View distance coefficient must be greater than 0" << std::endl;
        return false;
    }
    if (params.frameCount <= 0) {
        std::cout << "Error: Frame count must be greater than 0" << std::endl;
        return false;
    }

    // Ensure output directory exists
    std::filesystem::create_directories(params.outputDir);

    return true;
}

// Create orthographic projection camera
osg::ref_ptr<osg::Camera> createOrthoCamera(const osg::BoundingBox& bb, const osg::Vec3& direction,
    const RenderParameters& params)
{
    osg::ref_ptr<osg::Camera> camera = new osg::Camera;
    camera->setClearMask(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    camera->setClearColor(osg::Vec4(1.0f, 1.0f, 1.0f, 1.0f));
    camera->setRenderOrder(osg::Camera::PRE_RENDER);
    camera->setRenderTargetImplementation(osg::Camera::FRAME_BUFFER_OBJECT);
    camera->setReferenceFrame(osg::Transform::ABSOLUTE_RF);

    // Create texture
    osg::ref_ptr<osg::Texture2D> texture = new osg::Texture2D;
    texture->setTextureSize(params.imageWidth, params.imageHeight);
    texture->setInternalFormat(GL_RGB);
    texture->setFilter(osg::Texture2D::MIN_FILTER, osg::Texture2D::LINEAR);
    texture->setFilter(osg::Texture2D::MAG_FILTER, osg::Texture2D::LINEAR);

    // Create image
    osg::ref_ptr<osg::Image> image = new osg::Image;
    image->allocateImage(params.imageWidth, params.imageHeight, 1, GL_RGB, GL_UNSIGNED_BYTE);
    texture->setImage(image.get());

    camera->attach(osg::Camera::COLOR_BUFFER, image.get());

    // Calculate bounding box dimensions
    float width = bb.xMax() - bb.xMin();
    float height = bb.yMax() - bb.yMin();
    float depth = bb.zMax() - bb.zMin();
    float maxSize = std::max(std::max(width, height), depth);
    float margin = maxSize * params.marginRatio;

    // Calculate projection plane dimensions, maintaining aspect ratio
    float aspectRatio = static_cast<float>(params.imageWidth) / params.imageHeight;
    float projWidth, projHeight;

    if (direction.x() != 0) {
        projWidth = height;
        projHeight = depth;
    }
    else if (direction.y() != 0) {
        projWidth = width;
        projHeight = depth;
    }
    else {
        projWidth = width;
        projHeight = height;
    }

    // Adjust dimensions to maintain aspect ratio
    float currentAspect = projWidth / projHeight;
    if (currentAspect > aspectRatio) {
        projHeight = projWidth / aspectRatio;
    }
    else {
        projWidth = projHeight * aspectRatio;
    }

    // Add margins
    projWidth += margin * 2;
    projHeight += margin * 2;

    // Set orthographic projection
    float halfWidth = projWidth * 0.5f;
    float halfHeight = projHeight * 0.5f;

    if (direction.x() != 0) {
        camera->setProjectionMatrixAsOrtho(-halfWidth, halfWidth,
            -halfHeight, halfHeight,
            -maxSize, maxSize);
    }
    else if (direction.y() != 0) {
        camera->setProjectionMatrixAsOrtho(-halfWidth, halfWidth,
            -halfHeight, halfHeight,
            -maxSize, maxSize);
    }
    else {
        camera->setProjectionMatrixAsOrtho(-halfWidth, halfWidth,
            -halfHeight, halfHeight,
            -maxSize, maxSize);
    }

    // Set view matrix
    osg::Vec3 center = bb.center();
    osg::Vec3 eye = center + direction * maxSize * params.viewDistance;
    osg::Vec3 up(0.0f, 0.0f, 1.0f);

    if (direction.z() != 0) {
        up = osg::Vec3(0.0f, 1.0f, 0.0f);
    }

    // Create rotation matrix for correct orientation
    osg::Matrix rotation;
    if (direction.x() > 0) {
        rotation.makeRotate(osg::PI_2, osg::Z_AXIS);
    }
    else if (direction.x() < 0) {
        rotation.makeRotate(-osg::PI_2, osg::Z_AXIS);
    }
    else if (direction.y() > 0) {
        rotation.makeRotate(osg::PI, osg::Z_AXIS);
    }
    else if (direction.z() < 0) {
        rotation.makeRotate(osg::PI, osg::X_AXIS);
    }

    osg::Vec3 rotatedUp = up * rotation;
    camera->setViewMatrixAsLookAt(eye, center, rotatedUp);

    return camera;
}

// Setup render state
void setupRenderState(osg::Node* node)
{
    // 递归遍历所有节点并设置状态
    class StateSetVisitor : public osg::NodeVisitor
    {
    public:
        StateSetVisitor() : osg::NodeVisitor(osg::NodeVisitor::TRAVERSE_ALL_CHILDREN) {}

        virtual void apply(osg::Node& node)
        {
            osg::StateSet* stateSet = node.getOrCreateStateSet();

            // 禁用所有纹理单元
            for (unsigned int i = 0; i < 32; ++i) {
                stateSet->setTextureMode(i, GL_TEXTURE_2D, osg::StateAttribute::OFF | osg::StateAttribute::OVERRIDE | osg::StateAttribute::PROTECTED);
            }

            // 设置纯黑色材质
            osg::ref_ptr<osg::Material> material = new osg::Material;
            material->setColorMode(osg::Material::AMBIENT_AND_DIFFUSE);
            material->setAmbient(osg::Material::FRONT_AND_BACK, osg::Vec4(0.0f, 0.0f, 0.0f, 1.0f));
            material->setDiffuse(osg::Material::FRONT_AND_BACK, osg::Vec4(0.0f, 0.0f, 0.0f, 1.0f));
            material->setSpecular(osg::Material::FRONT_AND_BACK, osg::Vec4(0.0f, 0.0f, 0.0f, 1.0f));
            material->setEmission(osg::Material::FRONT_AND_BACK, osg::Vec4(0.0f, 0.0f, 0.0f, 1.0f));
            stateSet->setAttributeAndModes(material.get(), osg::StateAttribute::ON | osg::StateAttribute::OVERRIDE | osg::StateAttribute::PROTECTED);

            // 完全禁用光照
            stateSet->setMode(GL_LIGHTING, osg::StateAttribute::OFF | osg::StateAttribute::OVERRIDE | osg::StateAttribute::PROTECTED);

            // 强制使用填充模式
            osg::ref_ptr<osg::PolygonMode> polygonMode = new osg::PolygonMode;
            polygonMode->setMode(osg::PolygonMode::FRONT_AND_BACK, osg::PolygonMode::FILL);
            stateSet->setAttributeAndModes(polygonMode.get(), osg::StateAttribute::ON | osg::StateAttribute::OVERRIDE | osg::StateAttribute::PROTECTED);

            // 启用深度测试
            stateSet->setMode(GL_DEPTH_TEST, osg::StateAttribute::ON | osg::StateAttribute::OVERRIDE | osg::StateAttribute::PROTECTED);

            traverse(node);
        }
    };

    // 应用访问器
    StateSetVisitor visitor;
    node->accept(visitor);
}

void captureProjections(osg::Node* node, const RenderParameters& params)
{
    // Calculate bounding box
    osg::ComputeBoundsVisitor boundVisitor;
    node->accept(boundVisitor);
    osg::BoundingBox boundingBox = boundVisitor.getBoundingBox();

    // Output bounding box information
    std::cout << "Model Bounding Box Information:" << std::endl;
    std::cout << "Center: " << boundingBox.center().x() << ", "
        << boundingBox.center().y() << ", "
        << boundingBox.center().z() << std::endl;
    std::cout << "Dimensions: " << std::endl;
    std::cout << "X: " << boundingBox.xMax() - boundingBox.xMin() << std::endl;
    std::cout << "Y: " << boundingBox.yMax() - boundingBox.yMin() << std::endl;
    std::cout << "Z: " << boundingBox.zMax() - boundingBox.zMin() << std::endl;

    // Setup render state
    setupRenderState(node);

    // Create graphics context for off-screen rendering
    osg::ref_ptr<osg::GraphicsContext::Traits> traits = new osg::GraphicsContext::Traits;
    traits->x = 0;
    traits->y = 0;
    traits->width = params.imageWidth;
    traits->height = params.imageHeight;
    traits->windowDecoration = false;
    traits->doubleBuffer = true;
    traits->pbuffer = true;
    traits->alpha = 8;
    traits->depth = 24;

    osg::ref_ptr<osg::GraphicsContext> gc = osg::GraphicsContext::createGraphicsContext(traits.get());
    if (!gc.valid())
    {
        std::cout << "Failed to create graphics context!" << std::endl;
        return;
    }

    // Define six directions
    osg::Vec3 directions[6] = {
        osg::Vec3(1.0, 0.0, 0.0),   // Right view
        osg::Vec3(-1.0, 0.0, 0.0),  // Left view
        osg::Vec3(0.0, 1.0, 0.0),   // Front view
        osg::Vec3(0.0, -1.0, 0.0),  // Back view
        osg::Vec3(0.0, 0.0, 1.0),   // Top view
        osg::Vec3(0.0, 0.0, -1.0)   // Bottom view
    };

    const char* viewNames[6] = { "right", "left", "front", "back", "top", "bottom" };

    // Create viewer
    osg::ref_ptr<osgViewer::Viewer> viewer = new osgViewer::Viewer;
    viewer->setThreadingModel(osgViewer::ViewerBase::SingleThreaded);

    // Create camera and capture projections for each direction
    for (int i = 0; i < 6; ++i)
    {
        osg::ref_ptr<osg::Camera> camera = createOrthoCamera(boundingBox, directions[i], params);
        camera->addChild(node);

        viewer->setSceneData(camera.get());
        viewer->getCamera()->setGraphicsContext(gc.get());
        viewer->getCamera()->setViewport(0, 0, params.imageWidth, params.imageHeight);
        viewer->realize();

        // Render multiple frames to ensure proper rendering
        for (int j = 0; j < params.frameCount; ++j) {
            viewer->frame();
        }

        // Save image
        osg::Image* image = dynamic_cast<osg::Image*>(camera->getBufferAttachmentMap()
            .find(osg::Camera::COLOR_BUFFER)->second._image.get());

        if (image)
        {
            std::string filename = params.outputDir + "/" + viewNames[i] + "_projection.png";
            osgDB::writeImageFile(*image, filename);
            std::cout << "Saved " << filename << std::endl;
        }
    }
}

int main(int argc, char** argv)
{
    RenderParameters params;
    std::string inputFile;

    if (!parseCommandLine(argc, argv, inputFile, params))
    {
        return 1;
    }

    // Read OSGB file
    osg::ref_ptr<osg::Node> node = osgDB::readNodeFile(inputFile);
    if (!node)
    {
        std::cout << "Failed to read file: " << inputFile << std::endl;
        return 1;
    }

    // Optimize scene graph
    osgUtil::Optimizer optimizer;
    optimizer.optimize(node.get());

    // Capture projections from six directions
    captureProjections(node.get(), params);

    return 0;
}